"""
智能旅游助手Agent
基于LangGraph实现，支持旅游攻略查询和路线规划功能
"""

import logging
import json
import uuid
from typing import Dict, List, Any, Optional, TypedDict, Annotated
from datetime import datetime
from enum import Enum

from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_community.llms import Tongyi

from tools.rag_tool import RAGTool
from tools.amap_tool import AmapTool, TravelMode
from utils.config import Config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class QueryType(Enum):
    """查询类型枚举"""
    TRAVEL_GUIDE = "travel_guide"  # 旅游攻略查询
    ROUTE_PLANNING = "route_planning"  # 路线规划
    UNKNOWN = "unknown"  # 未知类型

class AgentState(TypedDict):
    """Agent状态定义"""
    messages: Annotated[List[BaseMessage], add_messages]
    user_id: str
    session_id: str
    query_type: Optional[QueryType]
    user_intent: Optional[str]
    route_info: Optional[Dict[str, Any]]
    rag_result: Optional[Dict[str, Any]]
    route_result: Optional[Dict[str, Any]]
    relevance_check: Optional[Dict[str, Any]]
    final_answer: Optional[str]
    needs_clarification: bool
    clarification_question: Optional[str]

class TravelAgent:
    """智能旅游助手Agent"""
    
    def __init__(self):
        """初始化Agent"""
        # 验证配置
        Config.validate_config()
        
        # 初始化工具
        self.rag_tool = RAGTool()
        self.amap_tool = AmapTool()
        
        # 初始化LLM
        self.llm = Tongyi(
            model_name=Config.LLM_MODEL,
            dashscope_api_key=Config.DASHSCOPE_API_KEY,
            temperature=0.1,
            max_tokens=2000
        )
        
        # 用户会话存储
        self.user_sessions: Dict[str, Dict[str, Any]] = {}
        
        # 构建工作流图
        self.workflow = self._build_workflow()
        
        logger.info("智能旅游助手Agent初始化完成")
    
    def _build_workflow(self) -> StateGraph:
        """构建LangGraph工作流"""
        workflow = StateGraph(AgentState)
        
        # 添加节点
        workflow.add_node("classify_query", self._classify_query)
        workflow.add_node("handle_travel_guide", self._handle_travel_guide)
        workflow.add_node("extract_route_info", self._extract_route_info)
        workflow.add_node("plan_route", self._plan_route)
        workflow.add_node("check_relevance", self._check_relevance)
        workflow.add_node("generate_final_answer", self._generate_final_answer)
        workflow.add_node("ask_clarification", self._ask_clarification)
        
        # 设置入口点
        workflow.set_entry_point("classify_query")
        
        # 添加条件边
        workflow.add_conditional_edges(
            "classify_query",
            self._route_after_classification,
            {
                "travel_guide": "handle_travel_guide",
                "route_planning": "extract_route_info",
                "unknown": "check_relevance"
            }
        )
        
        workflow.add_conditional_edges(
            "extract_route_info",
            self._route_after_extraction,
            {
                "complete": "plan_route",
                "incomplete": "ask_clarification"
            }
        )
        
        workflow.add_edge("handle_travel_guide", "check_relevance")
        workflow.add_edge("plan_route", "check_relevance")
        workflow.add_edge("ask_clarification", END)
        workflow.add_edge("check_relevance", "generate_final_answer")
        workflow.add_edge("generate_final_answer", END)
        
        return workflow.compile()
    
    def _classify_query(self, state: AgentState) -> AgentState:
        """分类用户查询"""
        logger.info("正在分类用户查询...")
        
        last_message = state["messages"][-1].content
        user_id = state["user_id"]
        
        # 使用LLM分类查询类型
        classification_prompt = f"""
        请分析以下用户查询，判断查询类型：
        
        用户查询：{last_message}
        
        查询类型说明：
        1. travel_guide: 旅游攻略相关查询，如景点推荐、美食推荐、交通信息、休闲娱乐等
        2. route_planning: 路线规划相关查询，如"怎么去"、"路线"、"导航"等
        3. unknown: 无法明确分类的查询
        
        请只返回查询类型（travel_guide、route_planning或unknown），不要其他内容。
        """
        
        try:
            response = self.llm.invoke(classification_prompt)
            query_type_str = response.strip().lower()
            
            if "travel_guide" in query_type_str:
                query_type = QueryType.TRAVEL_GUIDE
            elif "route_planning" in query_type_str:
                query_type = QueryType.ROUTE_PLANNING
            else:
                query_type = QueryType.UNKNOWN
                
        except Exception as e:
            logger.error(f"查询分类失败: {str(e)}")
            query_type = QueryType.UNKNOWN
        
        state["query_type"] = query_type
        state["user_intent"] = last_message
        
        logger.info(f"查询分类结果: {query_type.value}")
        return state
    
    def _handle_travel_guide(self, state: AgentState) -> AgentState:
        """处理旅游攻略查询"""
        logger.info("正在处理旅游攻略查询...")
        
        question = state["user_intent"]
        
        try:
            # 调用RAG工具
            rag_result = self.rag_tool.answer_question(question)
            state["rag_result"] = rag_result
            
            logger.info("旅游攻略查询处理完成")
            
        except Exception as e:
            logger.error(f"处理旅游攻略查询时出错: {str(e)}")
            state["rag_result"] = {
                "question": question,
                "answer": f"抱歉，处理您的旅游攻略查询时出现错误: {str(e)}",
                "source_documents": []
            }
        
        return state
    
    def _extract_route_info(self, state: AgentState) -> AgentState:
        """提取路线规划信息"""
        logger.info("正在提取路线规划信息...")
        
        user_intent = state["user_intent"]
        
        # 使用LLM提取路线信息
        extraction_prompt = f"""
        请从以下用户查询中提取路线规划信息：
        
        用户查询：{user_intent}
        
        请提取以下信息：
        1. 出发地（如果明确提到）
        2. 目的地（如果明确提到）
        3. 出行方式偏好（如果明确提到）
        
        如果信息不完整，请说明缺少什么信息。
        
        请以JSON格式返回：
        {{
            "origin": "出发地或null",
            "destination": "目的地或null", 
            "travel_mode": "出行方式或null",
            "complete": true/false,
            "missing_info": "缺少的信息"
        }}
        """
        
        try:
            response = self.llm.invoke(extraction_prompt)
            
            # 尝试解析JSON响应
            try:
                route_info = json.loads(response)
            except json.JSONDecodeError:
                # 如果JSON解析失败，尝试从文本中提取信息
                route_info = self._parse_route_info_from_text(response, user_intent)
            
            state["route_info"] = route_info
            
            logger.info(f"路线信息提取结果: {route_info}")
            
        except Exception as e:
            logger.error(f"提取路线信息时出错: {str(e)}")
            state["route_info"] = {
                "origin": None,
                "destination": None,
                "travel_mode": None,
                "complete": False,
                "missing_info": "无法解析路线信息"
            }
        
        return state
    
    def _parse_route_info_from_text(self, text: str, user_intent: str) -> Dict[str, Any]:
        """从文本中解析路线信息"""
        route_info = {
            "origin": None,
            "destination": None,
            "travel_mode": None,
            "complete": False,
            "missing_info": ""
        }
        
        # 简单的关键词匹配
        text_lower = text.lower()
        user_intent_lower = user_intent.lower()
        
        # 检查是否包含出发地和目的地
        if "出发地" in text_lower or "起点" in text_lower:
            route_info["origin"] = "需要用户提供"
        if "目的地" in text_lower or "终点" in text_lower:
            route_info["destination"] = "需要用户提供"
        
        # 检查出行方式
        if "步行" in user_intent_lower or "走路" in user_intent_lower:
            route_info["travel_mode"] = "walking"
        elif "骑车" in user_intent_lower or "自行车" in user_intent_lower:
            route_info["travel_mode"] = "bicycling"
        elif "公交" in user_intent_lower or "地铁" in user_intent_lower or "公共交通" in user_intent_lower:
            route_info["travel_mode"] = "transit"
        elif "开车" in user_intent_lower or "自驾" in user_intent_lower:
            route_info["travel_mode"] = "driving"
        
        # 判断信息是否完整
        if route_info["origin"] and route_info["destination"]:
            route_info["complete"] = True
        else:
            missing = []
            if not route_info["origin"]:
                missing.append("出发地")
            if not route_info["destination"]:
                missing.append("目的地")
            route_info["missing_info"] = "、".join(missing)
        
        return route_info
    
    def _plan_route(self, state: AgentState) -> AgentState:
        """规划路线"""
        logger.info("正在规划路线...")
        
        route_info = state["route_info"]
        
        try:
            origin = route_info["origin"]
            destination = route_info["destination"]
            travel_mode_str = route_info.get("travel_mode", "driving")
            
            # 转换出行方式
            try:
                travel_mode = TravelMode(travel_mode_str)
            except ValueError:
                travel_mode = TravelMode.DRIVING  # 默认自驾
            
            # 调用高德地图工具规划路线
            route_result = self.amap_tool.plan_route(origin, destination, travel_mode)
            state["route_result"] = route_result
            
            logger.info("路线规划完成")
            
        except Exception as e:
            logger.error(f"规划路线时出错: {str(e)}")
            state["route_result"] = {
                "error": f"路线规划失败: {str(e)}"
            }
        
        return state
    
    def _check_relevance(self, state: AgentState) -> AgentState:
        """检查答案相关性"""
        logger.info("正在检查答案相关性...")
        
        user_intent = state["user_intent"]
        query_type = state["query_type"]
        
        # 根据查询类型检查相关性
        if query_type == QueryType.TRAVEL_GUIDE and state.get("rag_result"):
            answer = state["rag_result"]["answer"]
            relevance_check = self._check_travel_guide_relevance(user_intent, answer)
        elif query_type == QueryType.ROUTE_PLANNING and state.get("route_result"):
            route_result = state["route_result"]
            relevance_check = self._check_route_relevance(user_intent, route_result)
        else:
            relevance_check = {
                "relevant": False,
                "reason": "无法确定查询类型或缺少结果"
            }
        
        state["relevance_check"] = relevance_check
        logger.info(f"相关性检查结果: {relevance_check}")
        
        return state
    
    def _check_travel_guide_relevance(self, user_intent: str, answer: str) -> Dict[str, Any]:
        """检查旅游攻略答案相关性"""
        relevance_prompt = f"""
        请评估以下旅游攻略答案与用户问题的相关性：
        
        用户问题：{user_intent}
        系统答案：{answer}
        
        请判断答案是否与用户问题高度相关。如果答案包含"根据提供的信息无法回答该问题"或类似表述，则认为不相关。
        
        请以JSON格式返回：
        {{
            "relevant": true/false,
            "reason": "相关性判断理由"
        }}
        """
        
        try:
            response = self.llm.invoke(relevance_prompt)
            return json.loads(response)
        except Exception as e:
            logger.error(f"检查旅游攻略相关性时出错: {str(e)}")
            return {
                "relevant": "根据提供的信息无法回答该问题" not in answer,
                "reason": f"检查过程出错: {str(e)}"
            }
    
    def _check_route_relevance(self, user_intent: str, route_result: Dict[str, Any]) -> Dict[str, Any]:
        """检查路线规划答案相关性"""
        if "error" in route_result:
            return {
                "relevant": False,
                "reason": f"路线规划失败: {route_result['error']}"
            }
        
        return {
            "relevant": True,
            "reason": "路线规划成功"
        }
    
    def _generate_final_answer(self, state: AgentState) -> AgentState:
        """生成最终答案"""
        logger.info("正在生成最终答案...")
        
        user_intent = state["user_intent"]
        query_type = state["query_type"]
        relevance_check = state.get("relevance_check", {})
        
        if not relevance_check.get("relevant", False):
            final_answer = f"抱歉，我无法回答您的问题：{user_intent}。{relevance_check.get('reason', '')}"
        else:
            if (query_type == QueryType.TRAVEL_GUIDE or (hasattr(query_type, 'value') and query_type.value == "travel_guide")) and state.get("rag_result"):
                final_answer = state["rag_result"]["answer"]
            elif (query_type == QueryType.ROUTE_PLANNING or (hasattr(query_type, 'value') and query_type.value == "route_planning")) and state.get("route_result"):
                final_answer = self._format_route_result(state["route_result"])
            else:
                final_answer = "抱歉，我无法处理您的请求。"
        
        state["final_answer"] = final_answer
        
        # 添加AI消息到对话历史
        state["messages"].append(AIMessage(content=final_answer))
        
        logger.info("最终答案生成完成")
        return state
    
    def _format_route_result(self, route_result: Dict[str, Any]) -> str:
        """格式化路线规划结果"""
        if "error" in route_result:
            return f"路线规划失败：{route_result['error']}"
        
        mode = route_result.get("mode", "")
        origin = route_result.get("origin", {}).get("address", "")
        destination = route_result.get("destination", {}).get("address", "")
        
        result_lines = [
            f"📍 路线规划结果",
            f"出行方式：{self.amap_tool._get_mode_display_name(mode)}",
            f"起点：{origin}",
            f"终点：{destination}",
            ""
        ]
        
        if mode in ["walking", "bicycling"]:
            summary = route_result.get("summary", {})
            result_lines.extend([
                f"总距离：{summary.get('total_distance', '未知')}",
                f"预计时间：{summary.get('total_duration', '未知')}"
            ])
        elif mode == "transit":
            summary = route_result.get("summary", {})
            result_lines.extend([
                f"预计时间：{summary.get('total_duration', '未知')}",
                f"费用：{summary.get('total_cost', '未知')}",
                f"步行距离：{summary.get('walking_distance', '未知')}"
            ])
        elif mode == "driving":
            summary = route_result.get("summary", {})
            result_lines.extend([
                f"总距离：{summary.get('total_distance', '未知')}",
                f"预计时间：{summary.get('total_duration', '未知')}",
                f"过路费：{summary.get('tolls', '未知')}"
            ])
        
        return "\n".join(result_lines)
    
    def _ask_clarification(self, state: AgentState) -> AgentState:
        """询问用户澄清信息"""
        logger.info("正在询问用户澄清信息...")
        
        route_info = state["route_info"]
        missing_info = route_info.get("missing_info", "")
        
        clarification_question = f"为了为您规划最佳路线，我需要了解以下信息：{missing_info}。请提供这些信息。"
        
        state["needs_clarification"] = True
        state["clarification_question"] = clarification_question
        state["final_answer"] = clarification_question
        
        # 添加AI消息到对话历史
        state["messages"].append(AIMessage(content=clarification_question))
        
        logger.info("澄清问题已生成")
        return state
    
    def _route_after_classification(self, state: AgentState) -> str:
        """分类后的路由决策"""
        query_type = state["query_type"]
        return query_type.value
    
    def _route_after_extraction(self, state: AgentState) -> str:
        """信息提取后的路由决策"""
        route_info = state["route_info"]
        if route_info.get("complete", False):
            return "complete"
        else:
            return "incomplete"
    
    def get_or_create_session(self, user_id: str) -> str:
        """获取或创建用户会话"""
        if user_id not in self.user_sessions:
            self.user_sessions[user_id] = {
                "session_id": str(uuid.uuid4()),
                "created_at": datetime.now(),
                "messages": []
            }
        
        return self.user_sessions[user_id]["session_id"]
    
    def chat(self, user_id: str, message: str) -> str:
        """处理用户消息"""
        logger.info(f"处理用户 {user_id} 的消息: {message}")
        
        # 获取或创建会话
        session_id = self.get_or_create_session(user_id)
        
        # 获取会话历史
        session = self.user_sessions[user_id]
        messages = session["messages"]
        
        # 添加用户消息
        messages.append(HumanMessage(content=message))
        
        # 创建初始状态
        initial_state = AgentState(
            messages=messages,
            user_id=user_id,
            session_id=session_id,
            query_type=None,
            user_intent=None,
            route_info=None,
            rag_result=None,
            route_result=None,
            final_answer=None,
            needs_clarification=False,
            clarification_question=None
        )
        
        # 运行工作流
        try:
            final_state = self.workflow.invoke(initial_state)
            
            # 更新会话历史
            session["messages"] = final_state["messages"]
            
            return final_state["final_answer"]
            
        except Exception as e:
            logger.error(f"处理用户消息时出错: {str(e)}")
            error_message = f"抱歉，处理您的消息时出现错误: {str(e)}"
            messages.append(AIMessage(content=error_message))
            return error_message
    
    def get_session_history(self, user_id: str) -> List[Dict[str, Any]]:
        """获取用户会话历史"""
        if user_id not in self.user_sessions:
            return []
        
        session = self.user_sessions[user_id]
        history = []
        
        for message in session["messages"]:
            history.append({
                "type": "user" if isinstance(message, HumanMessage) else "assistant",
                "content": message.content,
                "timestamp": datetime.now().isoformat()
            })
        
        return history
    
    def clear_session(self, user_id: str) -> bool:
        """清除用户会话"""
        if user_id in self.user_sessions:
            del self.user_sessions[user_id]
            return True
        return False
